package com.gitee.feizns.explore.data_structure.tree.bs.mid;

import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.Set;

/**
 * @author feizns
 * @since 2020/3/13
 */
public class DistanceK {

    public static void main(String[] args) {
//        distanceK()
    }

    static List<Integer> ret = new LinkedList<>();

    static Map<TreeNode, TreeNode> parents = new HashMap<>();

    public static List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        if ( K > 0 ) {
            createParents(root, null);
            Set<TreeNode> visited = new HashSet<>();
            Queue<TreeNode> queue = new LinkedList<>();
            queue.add(target);
            visited.add(target);
            int level = 0;
            while ( !queue.isEmpty() ) {
                if ( level == K ) {
                    queue.forEach(n -> ret.add(n.val));
                    return ret;
                } else {
                    int size = queue.size();
                    for (int i = 0; i < size; i++) {
                        TreeNode node = queue.remove();
                        if ( node.left != null && !visited.contains(node.left) ) {
                            queue.add(node.left);
                            visited.add(node.left);
                        }
                        if ( node.right != null && !visited.contains(node.right) ) {
                            queue.add(node.right);
                            visited.add(node.right);
                        }
                        TreeNode parent = parents.get(node);
                        if ( parent != null && !visited.contains(parent) ) {
                            queue.add(parent);
                            visited.add(parent);
                        }
                    }
                }
            }
        } else
            ret.add(target.val);
        return ret;
    }

    private static void createParents(TreeNode curt, TreeNode parent) {
        if ( curt != null ) {
            parents.put(curt, parent);
            createParents(curt.left, curt);
            createParents(curt.right, curt);
        }
    }

}
